# huffmanCode.py

import heapq
import numpy as np



"""
def printHuffTree(huffTree, prefix = ''):
   if len(huffTree) == 2:
      print huffTree[1], prefix

   else:
      printHuffTree(huffTree[1], prefix + '0')
      printHuffTree(huffTree[2], prefix + '1')


def getSymbol(huffTree, huffmanCode):

   if len(huffmanCode) == 0:
      return huffTree[1]

   if (huffmanCode[0] == '0'):
      return getSymbol(huffTree[1], huffmanCode[1:])
   else:
      return getSymbol(huffTree[2], huffmanCode[1:])
"""

# input dictionary has a list of (symbol, freq)
# populates exampleData
def dicToFreqData (dicData):
   data = []
   n = np.size(dicData)/2
   for i in np.arange(n):
      symbol = dicData[i][0]
      freq = dicData[i][1]
      data.append((freq, symbol))
   return data




class Huffman:
   def __init__(self, dicData):
      self.dicData = dicData
      self.freqData = dicToFreqData(dicData)
      self.huffTree = self.makeHuffTree(self.freqData)
      self.Symbol2Code = {}
      self.Code2Symbol = {}
      self.makeHuffTables(self.huffTree)

   def makeHuffTree(self, symbolTupleList):
      trees = list(symbolTupleList)
      heapq.heapify(trees)
      while len(trees) > 1:
         childR, childL = heapq.heappop(trees), heapq.heappop(trees)
  
         parent = (childL[0] + childR[0], childL, childR)
         heapq.heappush(trees, parent)

      return trees[0]


   def makeHuffTables(self, huffTree, prefix = ''):
      if len(huffTree) == 2:
         self.Symbol2Code[huffTree[1]] = prefix
         #print " inserted in s2c. key= ", huffTree[1], " value= ", prefix
         self.Code2Symbol[prefix] = huffTree[1]
         #print " inserted in c2s. key= ", prefix, " value= ", huffTree[1]
      else:
         self.makeHuffTables(huffTree[1], prefix + '0')
         self.makeHuffTables(huffTree[2], prefix + '1')


if __name__ == '__main__':

   
   Data = [
     (0.124167  , 'e'),   
     (0.0969225 , 't'),   
     (0.0820011 , 'a'),   
     (0.0768052 , 'i'),   
     (0.0764055 , 'n'),   
     (0.0714095 , 'o'),   
     (0.0706768 , 's'),   
     (0.0668132 , 'r'),   
     (0.0448308 , 'l'),   
     (0.0363709 , 'd'),   
     (0.0350386 , 'h'),   
     (0.0344391 , 'c'),   
     (0.028777  , 'u'),   
     (0.0281775 , 'm'),   
     (0.0235145 , 'f'),   
     (0.0203171 , 'p'),   
     (0.0189182 , 'y'),   
     (0.0181188 , 'g'),   
     (0.0135225 , 'w'),   
     (0.0124567 , 'v'),   
     (0.0106581 , 'b'),   
     (0.00393019, 'k'),   
     (0.00219824, 'x'),   
     (0.0019984 , 'j'),   
     (0.0009325 , 'q'),   
     (0.000599  , 'z')   
   ]
   
   # huffTree is the table to used. Possibly create mulitiple tables
   huffTree = makeHuffTree(Data)

   Symbol2Code = {}
   Code2Symbol = {}
   makeHuffTables(huffTree)   # populates Symbol2Code and Code2Symbol dictionaries

   # printHuffTree(huffTree)      # prints the generated huffman tree
   # print Symbol2Code            # dictionary for symbol to code lookup
   # print Code2Symbol            # dictionary for code to symbol lookup

   symbol = 'h'
   code = Symbol2Code[symbol]
   retrievedSymbol = Code2Symbol[code]

   print "symbol: ", symbol, " code: ", code, " retrievedSymbol: ", retrievedSymbol



